# Copyright (C) 2024 Intel Corporation
# SPDX-License-Identifier: Apache-2.0

import asyncio
import os
import ast
import re
from typing import Dict, List, Optional
import radon.metrics as metrics
from radon.complexity import cc_visit
import black
import autopep8
import javalang
from yapf.yapflib.yapf_api import FormatCode
from google_java_format.google_java_format import Formatter

from comps import MegaServiceEndpoint, MicroService, ServiceOrchestrator, ServiceRoleType, ServiceType
from comps.cores.proto.api_protocol import (
    ChatCompletionRequest,
    ChatCompletionResponse,
    ChatCompletionResponseChoice,
    ChatMessage,
    UsageInfo,
)
from fastapi import Request
from fastapi.responses import StreamingResponse

MEGA_SERVICE_PORT = int(os.getenv("MEGA_SERVICE_PORT", 7777))
LLM_SERVICE_HOST_IP = os.getenv("LLM_SERVICE_HOST_IP", "0.0.0.0")
LLM_SERVICE_PORT = int(os.getenv("LLM_SERVICE_PORT", 9000))


class CodePreprocessor:
    def __init__(self):
        self.parsers = {
            "python": self._parse_python,
            "java": self._parse_java,
            # 后续可添加其他语言解析器
        }
        
    def analyze_code(self, source_code: str, language: str) -> dict:
        """分析代码结构,返回关键信息"""
        if language.lower() not in self.parsers:
            raise ValueError(f"Unsupported language: {language}")
            
        analysis = {
            "functions": [],
            "classes": [],
            "dependencies": [],
            "complexity": 0,
            "lines_of_code": len(source_code.splitlines()),
            "syntax_valid": False
        }
        
        # 验证语法
        if not self.validate_syntax(source_code, language):
            return analysis
            
        analysis["syntax_valid"] = True
        
        # 解析代码结构
        parser_result = self.parsers[language.lower()](source_code)
        analysis.update(parser_result)
        
        # 计算复杂度
        analysis["complexity"] = self._calculate_complexity(source_code, language)
        
        return analysis
    
    def _parse_python(self, code: str) -> dict:
        """Python代码解析器"""
        result = {
            "functions": [],
            "classes": [],
            "dependencies": []
        }
        
        try:
            tree = ast.parse(code)
            
            # 提取导入
            for node in ast.walk(tree):
                if isinstance(node, ast.Import):
                    result["dependencies"].extend(n.name for n in node.names)
                elif isinstance(node, ast.ImportFrom):
                    result["dependencies"].append(node.module)
                    
            # 提取函数
            for node in ast.walk(tree):
                if isinstance(node, ast.FunctionDef):
                    func_info = {
                        "name": node.name,
                        "args": [arg.arg for arg in node.args.args],
                        "decorators": [ast.unparse(d) for d in node.decorator_list],
                        "line_number": node.lineno
                    }
                    result["functions"].append(func_info)
                    
            # 提取类
            for node in ast.walk(tree):
                if isinstance(node, ast.ClassDef):
                    class_info = {
                        "name": node.name,
                        "bases": [ast.unparse(b) for b in node.bases],
                        "methods": [m.name for m in node.body if isinstance(m, ast.FunctionDef)],
                        "line_number": node.lineno
                    }
                    result["classes"].append(class_info)
                    
        except SyntaxError as e:
            print(f"Python parsing error: {e}")
            
        return result
    
    def _parse_java(self, code: str) -> dict:
        """Java代码解析器"""
        result = {
            "functions": [],
            "classes": [],
            "dependencies": []
        }
        
        # 使用正则表达式提取基本信息
        # 导入语句
        imports = re.findall(r'import\s+([\w.]+);', code)
        result["dependencies"] = imports
        
        # 类定义
        class_pattern = r'class\s+(\w+)(?:\s+extends\s+(\w+))?(?:\s+implements\s+([\w,\s]+))?'
        for match in re.finditer(class_pattern, code):
            class_info = {
                "name": match.group(1),
                "parent": match.group(2),
                "interfaces": match.group(3).split(',') if match.group(3) else [],
                "line_number": code[:match.start()].count('\n') + 1
            }
            result["classes"].append(class_info)
            
        # 方法定义
        method_pattern = r'(?:public|private|protected)?\s+(?:static\s+)?(?:[\w<>[\],\s]+)\s+(\w+)\s*\([^)]*\)'
        for match in re.finditer(method_pattern, code):
            func_info = {
                "name": match.group(1),
                "line_number": code[:match.start()].count('\n') + 1
            }
            result["functions"].append(func_info)
            
        return result
    
    def validate_syntax(self, code: str, language: str) -> bool:
        """验证代码语法"""
        if language.lower() == "python":
            try:
                ast.parse(code)
                return True
            except SyntaxError:
                return False
        elif language.lower() == "java":
            # 简单的Java语法检查
            basic_checks = [
                r'\{.*\}',  # 检查大括号匹配
                r'class\s+\w+',  # 检查类定义
                r';$'  # 检查语句结尾分号
            ]
            return all(re.search(pattern, code, re.MULTILINE) for pattern in basic_checks)
        return False
    
    def _calculate_complexity(self, code: str, language: str) -> float:
        """计算代码复杂度"""
        if language.lower() == "python":
            try:
                # 使用radon计算圈复杂度
                complexity = max((item.complexity for item in cc_visit(code)), default=1)
                # 计算Halstead复杂度
                h = metrics.h_visit(code)
                # 综合评分 (1-10)
                return min(10, (complexity + h.total.h1 / 100) / 2)
            except:
                return 1.0
        else:
            # 其他语言使用简单估算
            indicators = [
                len(code.splitlines()),  # 代码行数
                code.count('if '),       # if语句数
                code.count('for '),      # 循环数
                code.count('while '),    # while循环数
                code.count('switch '),   # switch语句数
                code.count('catch '),    # 异常处理数
            ]
            return min(10, sum(indicators) / 50)  # 简单归一化到1-10


class LanguageConverter:
    def __init__(self):
        # 基础语法转换规则
        self.syntax_rules = {
            ("python", "java"): [
                # 函数定义转换
                (r"def\s+(\w+)\s*\((.*?)\):", r"public void $1($2) {"),
                # 类定义转换
                (r"class\s+(\w+):", r"public class $1 {"),
                # 构造函数转换
                (r"def\s+__init__\s*\(self,\s*(.*?)\):", r"public $CLASS_NAME($1) {"),
                # self引用转换
                (r"self\.(\w+)", r"this.$1"),
                # 列表定义转换
                (r"(\w+)\s*=\s*\[(.*?)\]", r"List<Object> $1 = Arrays.asList($2)"),
                # 字典定义转换
                (r"(\w+)\s*=\s*\{(.*?)\}", r"Map<String, Object> $1 = new HashMap<>() {{$2}}")
            ],
            ("java", "python"): [
                # 访问修饰符移除
                (r"(public|private|protected)\s+", r""),
                # 类型声明移除
                (r"(\w+)\s+(\w+)\s*=", r"$2 ="),
                # 构造函数转换
                (r"public\s+(\w+)\((.*?)\)\s*\{", r"def __init__(self, $2):"),
                # this引用转换
                (r"this\.(\w+)", r"self.$1"),
                # System.out.println转换
                (r"System\.out\.println\((.*?)\);", r"print($1)"),
                # 数组声明转换
                (r"(\w+)\[\]\s+(\w+)", r"$2: List[$1]")
            ]
        }
        
        # 常用API映射
        self.api_mappings = {
            ("python", "java"): {
                "len": "size()",
                "print": "System.out.println",
                "str": "String.valueOf",
                "list": "ArrayList",
                "dict": "HashMap",
                "set": "HashSet",
                "max": "Collections.max",
                "min": "Collections.min",
                "sum": "stream().mapToInt(Integer::valueOf).sum()",
            },
            ("java", "python"): {
                "String.valueOf": "str",
                "Integer.parseInt": "int",
                "System.out.println": "print",
                "ArrayList": "list",
                "HashMap": "dict",
                "HashSet": "set",
            }
        }
        
        # 导入语句映射
        self.import_mappings = {
            ("python", "java"): {
                "collections": "java.util.*",
                "datetime": "java.time.*",
                "json": "org.json.*",
                "math": "java.lang.Math",
                "random": "java.util.Random",
                "re": "java.util.regex.*",
            },
            ("java", "python"): {
                "java.util": "collections",
                "java.time": "datetime",
                "java.lang.Math": "math",
                "java.util.Random": "random",
                "java.util.regex": "re",
            }
        }

    def apply_rules(self, code: str, from_lang: str, to_lang: str) -> str:
        """应用语言转换规则"""
        # 转换为小写以统一处理
        from_lang = from_lang.lower()
        to_lang = to_lang.lower()
        
        # 获取对应的规则集
        lang_pair = (from_lang, to_lang)
        if lang_pair not in self.syntax_rules:
            return code
            
        # 1. 处理导入语句
        code = self._convert_imports(code, lang_pair)
        
        # 2. 应用语法转换规则
        for pattern, replacement in self.syntax_rules[lang_pair]:
            code = re.sub(pattern, replacement, code, flags=re.MULTILINE)
            
        # 3. 应用API映射
        code = self._apply_api_mappings(code, lang_pair)
        
        # 4. 特殊处理
        code = self._apply_special_rules(code, lang_pair)
        
        return code
        
    def _convert_imports(self, code: str, lang_pair: tuple) -> str:
        """转换导入语句"""
        if lang_pair not in self.import_mappings:
            return code
            
        mappings = self.import_mappings[lang_pair]
        
        if lang_pair[0] == "python":
            # Python to Java
            for py_import, java_import in mappings.items():
                code = re.sub(
                    rf"import\s+{py_import}",
                    f"import {java_import};",
                    code
                )
        else:
            # Java to Python
            for java_import, py_import in mappings.items():
                code = re.sub(
                    rf"import\s+{java_import}.*;",
                    f"import {py_import}",
                    code
                )
                
        return code
        
    def _apply_api_mappings(self, code: str, lang_pair: tuple) -> str:
        """应用API映射"""
        if lang_pair not in self.api_mappings:
            return code
            
        mappings = self.api_mappings[lang_pair]
        for source_api, target_api in mappings.items():
            code = code.replace(source_api, target_api)
            
        return code
        
    def _apply_special_rules(self, code: str, lang_pair: tuple) -> str:
        """应用特殊规则"""
        if lang_pair == ("python", "java"):
            # 添加类型声明
            code = re.sub(
                r"(\w+)\s*=\s*(\d+)",
                r"int $1 = $2",
                code
            )
            # 添加分号
            code = re.sub(r"([^;])\n", r"\1;\n", code)
            
        elif lang_pair == ("java", "python"):
            # 移除分号
            code = re.sub(r";", "", code)
            # 调整缩进
            code = self._fix_python_indentation(code)
            
        return code
        
    def _fix_python_indentation(self, code: str) -> str:
        """修复Python代码的缩进"""
        lines = code.split("\n")
        result = []
        indent_level = 0
        
        for line in lines:
            # 减少缩进的情况
            if line.strip().startswith("}"):
                indent_level -= 1
            
            # 添加正确的缩进
            if line.strip():  # 非空行
                result.append("    " * indent_level + line.strip())
            else:
                result.append("")
                
            # 增加缩进的情况
            if line.strip().endswith("{"):
                indent_level += 1
                
        return "\n".join(result)


class CodePostProcessor:
    def __init__(self):
        self.formatters = {
            "python": self._format_python,
            "java": self._format_java
        }
        self.validators = {
            "python": self._validate_python,
            "java": self._validate_java
        }
        
    def format_code(self, code: str, language: str) -> str:
        """代码格式化处理"""
        if language.lower() not in self.formatters:
            return code
            
        try:
            return self.formatters[language.lower()](code)
        except Exception as e:
            print(f"Format error: {e}")
            return code
            
    def _format_python(self, code: str) -> str:
        """Python代码格式化"""
        try:
            # 首先尝试使用black
            formatted_code = black.format_str(code, mode=black.FileMode())
        except:
            try:
                # 如果black失败,尝试使用autopep8
                formatted_code = autopep8.fix_code(code)
            except:
                try:
                    # 最后尝试使用yapf
                    formatted_code, _ = FormatCode(code)
                except:
                    return code
                    
        return formatted_code
        
    def _format_java(self, code: str) -> str:
        """Java代码格式��"""
        try:
            # 使用google-java-format
            formatter = Formatter()
            return formatter.format_code(code)
        except:
            return code
            
    def add_comments(self, code: str, original_code: str, language: str) -> str:
        """添加解释性注释"""
        # 提取原始代码的注释
        original_comments = self._extract_comments(original_code, language)
        
        # 分析代码结构变化
        changes = self._analyze_changes(original_code, code, language)
        
        # 添加转换说明注释
        code = self._add_conversion_comments(code, changes, language)
        
        # 保留原始重要注释
        code = self._preserve_comments(code, original_comments, language)
        
        return code
        
    def _extract_comments(self, code: str, language: str) -> List[Dict]:
        """提取代码中的注释"""
        comments = []
        
        if language.lower() == "python":
            lines = code.split("\n")
            for i, line in enumerate(lines):
                # 提取单行注释
                if "#" in line:
                    comment_text = line[line.index("#"):].strip()
                    comments.append({
                        "type": "inline",
                        "text": comment_text,
                        "line": i + 1
                    })
                    
                # 提取文档字符串
                if '"""' in line or "'''" in line:
                    # TODO: 处理多行文档字符串
                    pass
                    
        elif language.lower() == "java":
            # 使用javalang解析Java注释
            try:
                tokens = list(javalang.tokenizer.tokenize(code))
                for token in tokens:
                    if isinstance(token, javalang.tokenizer.Comment):
                        comments.append({
                            "type": "block" if token.value.startswith("/*") else "inline",
                            "text": token.value,
                            "line": token.position[0]
                        })
            except:
                pass
                
        return comments
        
    def _analyze_changes(self, original: str, converted: str, language: str) -> List[Dict]:
        """分析代码转换前后的变化"""
        changes = []
        
        # 分析结构变化
        original_structure = self._get_code_structure(original, language)
        converted_structure = self._get_code_structure(converted, language)
        
        # 记录重要变化
        changes.extend(self._compare_structures(
            original_structure,
            converted_structure
        ))
        
        # 分析API使用变化
        api_changes = self._analyze_api_changes(original, converted, language)
        changes.extend(api_changes)
        
        return changes
        
    def _get_code_structure(self, code: str, language: str) -> Dict:
        """获取代码结构信息"""
        structure = {
            "classes": [],
            "functions": [],
            "imports": []
        }
        
        if language.lower() == "python":
            try:
                tree = ast.parse(code)
                for node in ast.walk(tree):
                    if isinstance(node, ast.ClassDef):
                        structure["classes"].append(node.name)
                    elif isinstance(node, ast.FunctionDef):
                        structure["functions"].append(node.name)
                    elif isinstance(node, (ast.Import, ast.ImportFrom)):
                        # 处理导入
                        pass
            except:
                pass
                
        elif language.lower() == "java":
            try:
                tree = javalang.parse.parse(code)
                # 提取Java代码结构
                for path, node in tree.filter(javalang.tree.ClassDeclaration):
                    structure["classes"].append(node.name)
                for path, node in tree.filter(javalang.tree.MethodDeclaration):
                    structure["functions"].append(node.name)
            except:
                pass
                
        return structure
        
    def _compare_structures(self, original: Dict, converted: Dict) -> List[Dict]:
        """比较代码结构变化"""
        changes = []
        
        # 比较类
        for cls in converted["classes"]:
            if cls not in original["classes"]:
                changes.append({
                    "type": "class",
                    "name": cls,
                    "change": "added/renamed"
                })
                
        # 比较函数
        for func in converted["functions"]:
            if func not in original["functions"]:
                changes.append({
                    "type": "function",
                    "name": func,
                    "change": "added/renamed"
                })
                
        return changes
        
    def _analyze_api_changes(self, original: str, converted: str, language: str) -> List[Dict]:
        """分析API使用变化"""
        changes = []
        
        # TODO: 实��API变化分析
        
        return changes
        
    def _add_conversion_comments(self, code: str, changes: List[Dict], language: str) -> str:
        """添加转换说明注释"""
        comment_symbol = "#" if language.lower() == "python" else "//"
        
        # 添加文件头注释
        header_comment = f"{comment_symbol} Converted from {changes['source_language']} to {language}\n"
        header_comment += f"{comment_symbol} Important changes:\n"
        
        for change in changes:
            header_comment += f"{comment_symbol} - {change['type']}: {change['name']} ({change['change']})\n"
            
        return header_comment + "\n" + code
        
    def _preserve_comments(self, code: str, original_comments: List[Dict], language: str) -> str:
        """保留原始重要注释"""
        # TODO: 实现注释保留逻辑
        return code
        
    def validate_result(self, code: str, language: str) -> Dict:
        """验证转换结果"""
        if language.lower() not in self.validators:
            return {"valid": False, "errors": ["Unsupported language"]}
            
        return self.validators[language.lower()](code)
        
    def _validate_python(self, code: str) -> Dict:
        """验证Python代码"""
        result = {
            "valid": True,
            "errors": [],
            "warnings": []
        }
        
        # 语法检查
        try:
            ast.parse(code)
        except SyntaxError as e:
            result["valid"] = False
            result["errors"].append(f"Syntax error: {str(e)}")
            
        # TODO: 添加更多验证规则
        
        return result
        
    def _validate_java(self, code: str) -> Dict:
        """验证Java代码"""
        result = {
            "valid": True,
            "errors": [],
            "warnings": []
        }
        
        # 语法检查
        try:
            javalang.parse.parse(code)
        except Exception as e:
            result["valid"] = False
            result["errors"].append(f"Syntax error: {str(e)}")
            
        # TODO: 添加更多验证规则
        
        return result


class CodeTransService:
    def __init__(self, host="0.0.0.0", port=8000):
        self.host = host
        self.port = port
        self.megaservice = ServiceOrchestrator()
        self.endpoint = str(MegaServiceEndpoint.CODE_TRANS)

    def add_remote_service(self):
        llm = MicroService(
            name="llm",
            host=LLM_SERVICE_HOST_IP,
            port=LLM_SERVICE_PORT,
            endpoint="/v1/chat/completions",
            use_remote_service=True,
        )
        self.megaservice.add(llm)

    async def handle_request(self, request: Request):
        data = await request.json()
        language_from = data["language_from"]
        language_to = data["language_to"]
        source_code = data["source_code"]
        prompt_template = """
            ### System: Please translate the following {language_from} codes into {language_to} codes.

            ### Original codes:
            '''{language_from}

            {source_code}

            '''

            ### Translated codes:
        """
        prompt = prompt_template.format(language_from=language_from, language_to=language_to, source_code=source_code)
        result_dict, runtime_graph = await self.megaservice.schedule(initial_inputs={"query": prompt})
        for node, response in result_dict.items():
            # Here it suppose the last microservice in the megaservice is LLM.
            if (
                isinstance(response, StreamingResponse)
                and node == list(self.megaservice.services.keys())[-1]
                and self.megaservice.services[node].service_type == ServiceType.LLM
            ):
                return response
        last_node = runtime_graph.all_leaves()[-1]
        response = result_dict[last_node]["text"]
        choices = []
        usage = UsageInfo()
        choices.append(
            ChatCompletionResponseChoice(
                index=0,
                message=ChatMessage(role="assistant", content=response),
                finish_reason="stop",
            )
        )
        return ChatCompletionResponse(model="codetrans", choices=choices, usage=usage)

    def start(self):
        self.service = MicroService(
            self.__class__.__name__,
            service_role=ServiceRoleType.MEGASERVICE,
            host=self.host,
            port=self.port,
            endpoint=self.endpoint,
            input_datatype=ChatCompletionRequest,
            output_datatype=ChatCompletionResponse,
        )
        self.service.add_route(self.endpoint, self.handle_request, methods=["POST"])
        self.service.start()


class EnhancedCodeTransService(CodeTransService):
    def __init__(self, host="0.0.0.0", port=8000):
        super().__init__(host, port)
        self.preprocessor = CodePreprocessor()
        self.converter = LanguageConverter()
        self.postprocessor = CodePostProcessor()
        
    async def handle_request(self, request: Request):
        data = await request.json()
        
        # 1. 预处理
        analysis = self.preprocessor.analyze_code(
            data["source_code"], 
            data["language_from"]
        )
        
        # 2. 构建增强prompt
        prompt = self.build_enhanced_prompt(data, analysis)
        
        # 3. LLM转换
        result = await self.get_llm_response(prompt)
        
        # 4. 应用特定规则
        result = self.converter.apply_rules(
            result,
            data["language_from"],
            data["language_to"]
        )
        
        # 5. 后处理
        result = self.postprocessor.format_code(result, data["language_to"])
        
        return self.build_response(result)


class PromptBuilder:
    def __init__(self):
        # 语言特性映射
        self.language_features = {
            "python": {
                "typing": "动态类型,可选类型注解",
                "paradigm": "面向对象,函数式编程",
                "memory": "自动内存管理",
                "features": ["列表推导式", "装饰器", "生成器", "异步支持"],
                "conventions": ["PEP8", "缩进语法", "snake_case命名"]
            },
            "java": {
                "typing": "静态类型,强类型",
                "paradigm": "面向对象",
                "memory": "垃圾回收",
                "features": ["泛型", "注解", "接口", "多线程"],
                "conventions": ["驼峰命名法", "花括号语法", "访问修饰符"]
            }
        }
        
        # 转换重点提示
        self.conversion_focus = {
            ("python", "java"): [
                "添加类型声明",
                "添加访问修饰符",
                "转换列表/字典操作",
                "处理异常语法差异",
                "转换lambda表达式",
                "处理包/模块导入"
            ],
            ("java", "python"): [
                "移除类型声明",
                "简化访问修饰符",
                "转换集合操作",
                "调整异常处理",
                "转换匿名类",
                "调整包导入"
            ]
        }

    def build_prompt(self, source_code: str, analysis: dict, from_lang: str, to_lang: str) -> str:
        """构建完整的转换提示"""
        prompt = self._build_system_context(from_lang, to_lang)
        prompt += self._build_code_analysis(analysis)
        prompt += self._build_conversion_focus(from_lang, to_lang)
        prompt += self._build_code_section(source_code, from_lang)
        prompt += self._build_instructions(from_lang, to_lang)
        return prompt
        
    def _build_system_context(self, from_lang: str, to_lang: str) -> str:
        """构建系统上下文"""
        return f"""### System: 你是一个专业的代码转换助手
作为一个精通 {from_lang} 和 {to_lang} 的专家，请将以下代码从 {from_lang} 转换为 {to_lang}。

源语言({from_lang})特性:
{self._format_language_features(from_lang)}

目标语言({to_lang})特性:
{self._format_language_features(to_lang)}

"""

    def _format_language_features(self, lang: str) -> str:
        """格式化语言特性描述"""
        features = self.language_features.get(lang.lower(), {})
        return f"""- 类型系统: {features.get('typing', '')}
- 编程范式: {features.get('paradigm', '')}
- 内存管理: {features.get('memory', '')}
- 特有功能: {', '.join(features.get('features', []))}
- 编码规范: {', '.join(features.get('conventions', []))}
"""

    def _build_code_analysis(self, analysis: dict) -> str:
        """构建代码分析部分"""
        return f"""### 代码分析:
- 复杂度: {analysis.get('complexity', 'N/A')}/10
- 代码行数: {analysis.get('lines_of_code', 0)}
- 函数数量: {len(analysis.get('functions', []))}
- 类数量: {len(analysis.get('classes', []))}
- 依赖项: {', '.join(analysis.get('dependencies', []))}

函数列表:
{self._format_functions(analysis.get('functions', []))}

类定义:
{self._format_classes(analysis.get('classes', []))}

"""

    def _format_functions(self, functions: List[dict]) -> str:
        """格式化函数信息"""
        if not functions:
            return "无函数定义"
            
        result = []
        for func in functions:
            args = ', '.join(func.get('args', []))
            decorators = ' '.join(f'@{d}' for d in func.get('decorators', []))
            result.append(f"- {func['name']}({args}) {decorators}")
        return '\n'.join(result)

    def _format_classes(self, classes: List[dict]) -> str:
        """格式化类信息"""
        if not classes:
            return "无类定义"
            
        result = []
        for cls in classes:
            bases = ', '.join(cls.get('bases', []))
            methods = ', '.join(cls.get('methods', []))
            result.append(f"- {cls['name']}({bases}) -> methods: {methods}")
        return '\n'.join(result)

    def _build_conversion_focus(self, from_lang: str, to_lang: str) -> str:
        """构建转换重点提示"""
        lang_pair = (from_lang.lower(), to_lang.lower())
        focus_points = self.conversion_focus.get(lang_pair, [])
        
        if not focus_points:
            return ""
            
        return f"""### 转换重点:
{''.join(f'- {point}\n' for point in focus_points)}
"""

    def _build_code_section(self, source_code: str, from_lang: str) -> str:
        """构建源代码部分"""
        return f"""### 源代码 ({from_lang}):
```{from_lang}
{source_code}
```

"""

    def _build_instructions(self, from_lang: str, to_lang: str) -> str:
        """构建具体转换指"""
        return f"""### 转换要求:
1. 保持代码功能完全一致
2. 遵循{to_lang}的最佳实践和编码规范
3. 保留有意义的注释，并转换为目标语言的注释格式
4. 确保转换后的代码具有良好的可读性和可维护性
5. 适当使用{to_lang}的特有功能优化代码
6. 处理好语言间的类型系统差异
7. 注意异常处理的差异转换
8. 确保命名规范符合目标语言习惯

### 请输出转换后的{to_lang}代码:
```{to_lang}
"""

    def customize_for_complexity(self, prompt: str, complexity: float) -> str:
        """根据代码复杂度调整提示"""
        if complexity > 7:
            prompt += "\n注意：输入代码复杂度较高，请特别注意：\n"
            prompt += "1. 分步骤转换复杂逻辑\n"
            prompt += "2. 添加详细注释说明\n"
            prompt += "3. 考虑适当重构提高可维护性\n"
        return prompt

    def add_examples(self, prompt: str, from_lang: str, to_lang: str) -> str:
        """添加相关示例（可选）"""
        # TODO: 维护一个示例库，根据代码特征选择合适的示例
        return prompt

    def add_test_requirements(self, prompt: str) -> str:
        """添加测试相关要求（可选）"""
        prompt += "\n### 测试建议:\n"
        prompt += "1. 确保转换后的代码可以通过基本的单元测试\n"
        prompt += "2. 考虑边界情况的处理\n"
        prompt += "3. 保持异常处理的一致性\n"
        return prompt


class TestGenerator:
    def __init__(self):
        self.test_templates = {
            "python": {
                "import": "import unittest\n\n",
                "class": "class Test{name}(unittest.TestCase):\n",
                "method": "    def test_{name}(self):\n",
                "assert": "        self.assertEqual({expected}, {actual})\n"
            },
            "java": {
                "import": "import org.junit.Test;\nimport static org.junit.Assert.*;\n\n",
                "class": "public class Test{name} {\n",
                "method": "    @Test\n    public void test{name}() {\n",
                "assert": "        assertEquals({expected}, {actual});\n"
            }
        }

    def generate_tests(self, original_code: str, translated_code: str,
                      from_lang: str, to_lang: str) -> str:
        """生成对应的测试用例"""
        # 分析原始代码和转换后的代码
        original_analysis = self._analyze_code(original_code, from_lang)
        translated_analysis = self._analyze_code(translated_code, to_lang)
        
        # 生成测试用例
        test_cases = self._generate_test_cases(
            original_analysis,
            translated_analysis,
            from_lang,
            to_lang
        )
        
        # 根据目标语言格式化测试代码
        return self._format_test_code(test_cases, to_lang)
        
    def _analyze_code(self, code: str, language: str) -> Dict:
        """分析代码结构和功能点"""
        analysis = {
            "functions": [],
            "classes": [],
            "variables": [],
            "return_values": [],
            "control_flows": []
        }
        
        if language.lower() == "python":
            try:
                tree = ast.parse(code)
                # 提取函数信息
                for node in ast.walk(tree):
                    if isinstance(node, ast.FunctionDef):
                        analysis["functions"].append({
                            "name": node.name,
                            "args": [arg.arg for arg in node.args.args],
                            "returns": self._analyze_return_values(node)
                        })
                    elif isinstance(node, ast.ClassDef):
                        analysis["classes"].append({
                            "name": node.name,
                            "methods": [m.name for m in node.body if isinstance(m, ast.FunctionDef)]
                        })
            except:
                pass
                
        elif language.lower() == "java":
            try:
                tree = javalang.parse.parse(code)
                # 提取Java代码结构
                for path, node in tree.filter(javalang.tree.MethodDeclaration):
                    analysis["functions"].append({
                        "name": node.name,
                        "args": [(param.type.name, param.name) for param in node.parameters],
                        "returns": node.return_type.name if node.return_type else None
                    })
            except:
                pass
                
        return analysis
        
    def _analyze_return_values(self, node: ast.AST) -> List:
        """分析函数返回值"""
        returns = []
        for n in ast.walk(node):
            if isinstance(n, ast.Return) and n.value:
                returns.append(ast.unparse(n.value))
        return returns
        
    def _generate_test_cases(self, original: Dict, translated: Dict,
                           from_lang: str, to_lang: str) -> List[Dict]:
        """生成测试用例"""
        test_cases = []
        
        # 为每个函数生成测试
        for func in translated["functions"]:
            test_case = {
                "name": func["name"],
                "tests": []
            }
            
            # 基本功能测试
            test_case["tests"].extend(self._generate_basic_tests(func))
            
            # 边界条件测试
            test_case["tests"].extend(self._generate_edge_cases(func))
            
            # 异常测试
            test_case["tests"].extend(self._generate_exception_tests(func))
            
            test_cases.append(test_case)
            
        return test_cases
        
    def _generate_basic_tests(self, func: Dict) -> List[Dict]:
        """生成基本功能测试"""
        tests = []
        
        # 根据参数类型生成典型输入
        for args in self._generate_typical_inputs(func["args"]):
            tests.append({
                "name": f"basic_{func['name']}",
                "inputs": args,
                "expected": "TODO: Define expected output"
            })
            
        return tests
        
    def _generate_edge_cases(self, func: Dict) -> List[Dict]:
        """生成边界条件测试"""
        tests = []
        
        # 空值测试
        tests.append({
            "name": f"null_{func['name']}",
            "inputs": ["null" for _ in func["args"]],
            "expected": "TODO: Define expected output"
        })
        
        # 极限值测试
        tests.append({
            "name": f"limit_{func['name']}",
            "inputs": self._generate_limit_values(func["args"]),
            "expected": "TODO: Define expected output"
        })
        
        return tests
        
    def _generate_exception_tests(self, func: Dict) -> List[Dict]:
        """生成异常测试"""
        tests = []
        
        # 类型错误测试
        tests.append({
            "name": f"type_error_{func['name']}",
            "inputs": ["wrong_type" for _ in func["args"]],
            "expected_exception": "TypeError"
        })
        
        return tests
        
    def _generate_typical_inputs(self, args: List) -> List:
        """生成典型输入值"""
        typical_values = {
            "int": [0, 1, -1, 100],
            "float": [0.0, 1.0, -1.0, 3.14],
            "str": ["", "test", "hello world"],
            "list": [[], [1, 2, 3], ["a", "b", "c"]],
            "dict": [{}, {"key": "value"}, {"a": 1, "b": 2}]
        }
        
        return [typical_values.get(arg, ["default"]) for arg in args]
        
    def _generate_limit_values(self, args: List) -> List:
        """生成极限值"""
        limit_values = {
            "int": [sys.maxsize, -sys.maxsize - 1],
            "float": [float('inf'), float('-inf')],
            "str": ["" * 1000],  # 长字符串
            "list": [[] * 1000],  # 大列表
            "dict": [{}]  # 空字典
        }
        
        return [limit_values.get(arg, ["default"]) for arg in args]
        
    def _format_test_code(self, test_cases: List[Dict], language: str) -> str:
        """格式化测试代码"""
        template = self.test_templates[language.lower()]
        code = template["import"]
        
        for test_case in test_cases:
            class_name = f"Test{test_case['name'].capitalize()}"
            code += template["class"].format(name=class_name)
            
            for test in test_case["tests"]:
                code += template["method"].format(name=test["name"])
                
                if "expected_exception" in test:
                    if language.lower() == "python":
                        code += f"        with self.assertRaises({test['expected_exception']}):\n"
                        code += f"            {test_case['name']}({', '.join(map(str, test['inputs']))})\n"
                else:
                    code += template["assert"].format(
                        expected=test["expected"],
                        actual=f"{test_case['name']}({', '.join(map(str, test['inputs']))})"
                    )
                    
                code += "    }\n" if language.lower() == "java" else "\n"
                
            code += "}\n" if language.lower() == "java" else "\n"
            
        return code


class CodeValidator:
    def __init__(self):
        self.validators = {
            "syntax": self._validate_syntax,
            "structure": self._validate_structure,
            "functionality": self._validate_functionality,
            "style": self._validate_style
        }

    def validate_translation(self, original: str, translated: str,
                           from_lang: str, to_lang: str) -> Dict:
        """验证转换结果的正确性"""
        results = {
            "valid": True,
            "syntax_valid": False,
            "structure_preserved": False,
            "functionality_equivalent": False,
            "style_compliant": False,
            "issues": [],
            "warnings": [],
            "metrics": {}
        }
        
        # 运行所有验证器
        for validator_name, validator_func in self.validators.items():
            try:
                validator_result = validator_func(
                    original, translated, from_lang, to_lang
                )
                results.update(validator_result)
            except Exception as e:
                results["issues"].append(f"{validator_name} validation error: {str(e)}")
                results["valid"] = False
                
        # 计算相似度指标
        results["metrics"]["similarity"] = self._calculate_similarity(
            original, translated
        )
        
        # 计算复杂度变化
        results["metrics"]["complexity_change"] = self._compare_complexity(
            original, translated, from_lang, to_lang
        )
        
        return results
        
    def _validate_syntax(self, original: str, translated: str,
                        from_lang: str, to_lang: str) -> Dict:
        """语法验证"""
        result = {"syntax_valid": True, "issues": []}
        
        # 验证目标代码语法
        if to_lang.lower() == "python":
            try:
                ast.parse(translated)
            except SyntaxError as e:
                result["syntax_valid"] = False
                result["issues"].append(f"Python syntax error: {str(e)}")
        elif to_lang.lower() == "java":
            try:
                javalang.parse.parse(translated)
            except Exception as e:
                result["syntax_valid"] = False
                result["issues"].append(f"Java syntax error: {str(e)}")
                
        return result
        
    def _validate_structure(self, original: str, translated: str,
                          from_lang: str, to_lang: str) -> Dict:
        """结构一致性验证"""
        result = {"structure_preserved": True, "issues": []}
        
        # 分析原始代码和转换后代码的结构
        original_structure = self._analyze_structure(original, from_lang)
        translated_structure = self._analyze_structure(translated, to_lang)
        
        # 比较函数数量
        if len(original_structure["functions"]) != len(translated_structure["functions"]):
            result["structure_preserved"] = False
            result["issues"].append("Function count mismatch")
            
        # 比较类数量
        if len(original_structure["classes"]) != len(translated_structure["classes"]):
            result["structure_preserved"] = False
            result["issues"].append("Class count mismatch")
            
        return result
        
    def _validate_functionality(self, original: str, translated: str,
                              from_lang: str, to_lang: str) -> Dict:
        """功能等价性验证"""
        result = {"functionality_equivalent": True, "issues": []}
        
        # 生成测试用例
        test_generator = TestGenerator()
        tests = test_generator.generate_tests(original, translated, from_lang, to_lang)
        
        # 运行测试
        try:
            # TODO: 实现测试执行逻辑
            pass
        except Exception as e:
            result["functionality_equivalent"] = False
            result["issues"].append(f"Test execution error: {str(e)}")
            
        return result
        
    def _validate_style(self, original: str, translated: str,
                       from_lang: str, to_lang: str) -> Dict:
        """代码风格验证"""
        result = {"style_compliant": True, "issues": []}
        
        if to_lang.lower() == "python":
            # 使用pylint检查Python代码风格
            try:
                # TODO: 实现pylint检查
                pass
            except Exception as e:
                result["style_compliant"] = False
                result["issues"].append(f"Python style check error: {str(e)}")
        elif to_lang.lower() == "java":
            # 使用checkstyle检查Java代码风格
            try:
                # TODO: 实现checkstyle检查
                pass
            except Exception as e:
                result["style_compliant"] = False
                result["issues"].append(f"Java style check error: {str(e)}")
                
        return result
        
    def _calculate_similarity(self, original: str, translated: str) -> float:
        """计算代码相似度"""
        # 使用difflib计算相似度
        matcher = difflib.SequenceMatcher(None, original, translated)
        return matcher.ratio()
        
    def _compare_complexity(self, original: str, translated: str,
                          from_lang: str, to_lang: str) -> float:
        """比较代码复杂度变化"""
        try:
            if from_lang.lower() == "python" and to_lang.lower() == "python":
                # 使用radon计算Python代码复杂度
                original_complexity = ComplexityVisitor.from_code(original).total_complexity
                translated_complexity = ComplexityVisitor.from_code(translated).total_complexity
                return translated_complexity - original_complexity
        except:
            pass
            
        return 0.0
        
    def _analyze_structure(self, code: str, language: str) -> Dict:
        """分析代码结构"""
        structure = {
            "functions": [],
            "classes": [],
            "imports": []
        }
        
        if language.lower() == "python":
            try:
                tree = ast.parse(code)
                for node in ast.walk(tree):
                    if isinstance(node, ast.FunctionDef):
                        structure["functions"].append(node.name)
                    elif isinstance(node, ast.ClassDef):
                        structure["classes"].append(node.name)
                    elif isinstance(node, (ast.Import, ast.ImportFrom)):
                        # 处理导入
                        pass
            except:
                pass
        elif language.lower() == "java":
            try:
                tree = javalang.parse.parse(code)
                for path, node in tree.filter(javalang.tree.MethodDeclaration):
                    structure["functions"].append(node.name)
                for path, node in tree.filter(javalang.tree.ClassDeclaration):
                    structure["classes"].append(node.name)
            except:
                pass
                
        return structure


class ModelSelector:
    def __init__(self):
        # 模型配置
        self.models = {
            "gpt-4": {
                "capabilities": {
                    "max_complexity": 10,
                    "supported_languages": {"python", "java", "c++", "javascript"},
                    "context_length": 8000,
                    "performance": 0.95
                },
                "cost": 0.1  # 每1k tokens的成本
            },
            "gpt-3.5-turbo": {
                "capabilities": {
                    "max_complexity": 7,
                    "supported_languages": {"python", "java", "javascript"},
                    "context_length": 4000,
                    "performance": 0.85
                },
                "cost": 0.02
            },
            "codellama-34b": {
                "capabilities": {
                    "max_complexity": 8,
                    "supported_languages": {"python", "java", "c++"},
                    "context_length": 6000,
                    "performance": 0.90
                },
                "cost": 0.01
            }
        }
        
        # 语言特定配置
        self.language_configs = {
            ("python", "java"): {
                "preferred_models": ["gpt-4", "codellama-34b"],
                "min_performance": 0.85
            },
            ("java", "python"): {
                "preferred_models": ["gpt-4", "codellama-34b"],
                "min_performance": 0.85
            }
        }

    def select_model(self, code_complexity: float, language_pair: tuple,
                    code_length: int = 0, budget: float = float('inf')) -> str:
        """根据代码特征选择最合适的模型"""
        from_lang, to_lang = language_pair
        suitable_models = []
        
        # 1. 筛选支持该语言对的模型
        for model_name, config in self.models.items():
            if (from_lang in config["capabilities"]["supported_languages"] and
                to_lang in config["capabilities"]["supported_languages"]):
                suitable_models.append(model_name)
                
        if not suitable_models:
            raise ValueError(f"No model supports {from_lang} to {to_lang} conversion")
            
        # 2. 根据代码复杂度筛选
        suitable_models = [
            model for model in suitable_models
            if self.models[model]["capabilities"]["max_complexity"] >= code_complexity
        ]
        
        # 3. 根据代码长度筛选
        if code_length > 0:
            suitable_models = [
                model for model in suitable_models
                if self.models[model]["capabilities"]["context_length"] >= code_length
            ]
            
        # 4. 根据预算筛选
        if budget < float('inf'):
            estimated_tokens = code_length * 2  # 粗略估算
            suitable_models = [
                model for model in suitable_models
                if (self.models[model]["cost"] * estimated_tokens / 1000) <= budget
            ]
            
        if not suitable_models:
            raise ValueError("No suitable model found with given constraints")
            
        # 5. 根据性能和成本平衡选择最佳模型
        best_model = max(
            suitable_models,
            key=lambda m: self.models[m]["capabilities"]["performance"] / 
                        self.models[m]["cost"]
        )
        
        return best_model
        
    def get_model_config(self, model_name: str) -> Dict:
        """获取模型配置"""
        return self.models.get(model_name, {})
        
    def estimate_cost(self, model_name: str, code_length: int) -> float:
        """估算转换成本"""
        if model_name not in self.models:
            return 0.0
            
        estimated_tokens = code_length * 2  # 粗略估算
        return self.models[model_name]["cost"] * estimated_tokens / 1000


class BatchProcessor:
    def __init__(self, service, max_workers: int = 5):
        self.service = service
        self.max_workers = max_workers
        self.model_selector = ModelSelector()
        self.executor = ThreadPoolExecutor(max_workers=max_workers)
        
    async def process_directory(self, dir_path: str, from_lang: str, to_lang: str,
                              recursive: bool = True) -> Dict:
        """处理整个目录的代码转换"""
        results = {
            "successful": [],
            "failed": [],
            "skipped": [],
            "total_files": 0,
            "success_rate": 0.0,
            "total_cost": 0.0,
            "processing_time": 0.0
        }
        
        start_time = asyncio.get_event_loop().time()
        
        # 获取所有需要处理的文件
        files = await self._get_source_files(dir_path, from_lang, recursive)
        results["total_files"] = len(files)
        
        # 创建输出目录
        output_dir = self._create_output_directory(dir_path, to_lang)
        
        # 并发处理文件
        tasks = []
        for file_path in files:
            task = self._process_file(file_path, output_dir, from_lang, to_lang)
            tasks.append(task)
            
        # 等待所有任务完成
        file_results = await asyncio.gather(*tasks, return_exceptions=True)
        
        # 统计结果
        for result in file_results:
            if isinstance(result, Exception):
                results["failed"].append(str(result))
            elif result["status"] == "success":
                results["successful"].append(result["file"])
                results["total_cost"] += result["cost"]
            elif result["status"] == "skipped":
                results["skipped"].append(result["file"])
                
        # 计算成功率和处理时间
        results["success_rate"] = len(results["successful"]) / results["total_files"]
        results["processing_time"] = asyncio.get_event_loop().time() - start_time
        
        return results
        
    async def _get_source_files(self, dir_path: str, language: str,
                               recursive: bool) -> List[str]:
        """获取需要处理的源文件"""
        extensions = {
            "python": {".py"},
            "java": {".java"},
            "javascript": {".js"},
            "c++": {".cpp", ".hpp", ".h"}
        }
        
        if language.lower() not in extensions:
            raise ValueError(f"Unsupported language: {language}")
            
        valid_extensions = extensions[language.lower()]
        files = []
        
        for ext in valid_extensions:
            if recursive:
                pattern = f"**/*{ext}"
            else:
                pattern = f"*{ext}"
                
            files.extend([
                str(p) for p in Path(dir_path).glob(pattern)
                if p.is_file()
            ])
            
        return files
        
    def _create_output_directory(self, source_dir: str, target_lang: str) -> str:
        """创建输出目录"""
        output_dir = os.path.join(
            source_dir,
            f"converted_{target_lang.lower()}"
        )
        os.makedirs(output_dir, exist_ok=True)
        return output_dir
        
    async def _process_file(self, file_path: str, output_dir: str,
                           from_lang: str, to_lang: str) -> Dict:
        """处理单个文件"""
        try:
            # 读取源文件
            async with aiofiles.open(file_path, 'r') as f:
                source_code = await f.read()
                
            # 分析代码复杂度
            code_complexity = self._analyze_complexity(source_code)
            
            # 选择合适的模型
            model = self.model_selector.select_model(
                code_complexity,
                (from_lang, to_lang),
                len(source_code)
            )
            
            # 执行转换
            translated_code = await self.service.translate_code(
                source_code,
                from_lang,
                to_lang,
                model
            )
            
            # 生成输出文件路径
            relative_path = os.path.relpath(file_path, start=os.path.dirname(output_dir))
            output_path = os.path.join(
                output_dir,
                self._convert_file_extension(relative_path, to_lang)
            )
            
            # 确保输出目录存在
            os.makedirs(os.path.dirname(output_path), exist_ok=True)
            
            # 写入转换后的代码
            async with aiofiles.open(output_path, 'w') as f:
                await f.write(translated_code)
                
            return {
                "status": "success",
                "file": file_path,
                "output": output_path,
                "model": model,
                "cost": self.model_selector.estimate_cost(model, len(source_code))
            }
            
        except Exception as e:
            return {
                "status": "failed",
                "file": file_path,
                "error": str(e)
            }
            
    def _analyze_complexity(self, code: str) -> float:
        """分析代码复杂度"""
        try:
            visitor = ComplexityVisitor.from_code(code)
            return visitor.total_complexity
        except:
            return 5.0  # 默认中等复杂度
            
    def _convert_file_extension(self, file_path: str, target_lang: str) -> str:
        """转换文件扩展名"""
        extensions = {
            "python": ".py",
            "java": ".java",
            "javascript": ".js",
            "c++": ".cpp"
        }
        
        base = os.path.splitext(file_path)[0]
        return base + extensions.get(target_lang.lower(), ".txt")


if __name__ == "__main__":
    service_ochestrator = CodeTransService(port=MEGA_SERVICE_PORT)
    service_ochestrator.add_remote_service()
    service_ochestrator.start()
